#include <math.h>
#include <stdlib.h>
#include <string.h>
#include <mpi.h>
#include <omp.h>
#include <limits.h>

#include <cassert>
#include <fstream>
#include <iostream>
#include <string>
#include <vector>
#include <numeric>
#include <algorithm>

#include "OptionParser.h"
#include "ResultDatabase.h"
#include "Timer.h"

#if !defined (IDEAL_ALIGNMENT)
/** Default alignment, in bytes. */
#  define IDEAL_ALIGNMENT 16
#endif


using namespace std;

template <class Key, class Val>
void RunTest(string testName, ResultDatabase &resultDB, OptionParser &op);

// ****************************************************************************
// Function: addBenchmarkSpecOptions
//
// Purpose:
//   Add benchmark specific options parsing
//
// Arguments:
//   op: the options parser / parameter database
//
// Returns:  nothing
//
// Programmer: Kyle Spafford
// Creation: August 13, 2009
//
// Modifications:
//
// ****************************************************************************
void
addBenchmarkSpecOptions(OptionParser &op)
{
    op.addOption("mb", OPT_INT, "0", "data size (in MB)");
    op.addOption("rw", OPT_INT, "16", "radix width in bits");
}

// ****************************************************************************
// Function: RunBenchmark
//
// Purpose:
//   Driver for the generic benchmark. Templated for testing in different 
//   precition.
//
// Arguments:
//   resultDB: results from the benchmark are stored in this db
//   op: the options parser / parameter database
//
// Returns:  nothing
//
// Programmer: Kyle Spafford
// Creation: August 13, 2009
//
// Modifications:
//
// ****************************************************************************
void
RunBenchmark(ResultDatabase &resultDB, OptionParser &op)
{
    RunTest<long long, long long>("long", resultDB, op);
}


// ****************************************************************************
// Function: RunTest
//
// Purpose:
//   Primary method for the benchmark
//
// Arguments:
//   testName: the name of the test currently being executed (specifying SP or
//             DP)
//   resultDB: results from the benchmark are stored in this db
//   op: the options parser / parameter database
//
// Returns:  nothing
//
// Programmer: Kyle Spafford
// Creation: August 13, 2009
//
// Modifications:
//
// ****************************************************************************
template <class Key, class Val>
void RunTest(string testName, ResultDatabase &resultDB, OptionParser &op)
{
    int mpi_size, mpi_rank;
	MPI_Comm_size(MPI_COMM_WORLD, &mpi_size);
	MPI_Comm_rank(MPI_COMM_WORLD, &mpi_rank);
    
    // Problem Size Initialization
    // Number of Keys\Vals to sort (in MB)
    int probSizes[4] = { 1, 8, 16, 32 };
    int size = probSizes[op.getOptionInt("size")-1];
    
    // Convert to MB
    size = (size * 1000 * 1000) / sizeof(Key);
    
    // If the user specified a custom size, use that instead
    if (op.getOptionInt("mb") != 0)
    {
        size = op.getOptionInt("mb") * 1000 * 1000 / sizeof(Key);
    }
	
    int num_bits = sizeof(Key) * 8;
    int num_radix_bits = op.getOptionInt("rw");
 
    // Number of possible digit combinations, given the number of radix bits
    int num_buckets = (int) pow (2.0, num_radix_bits);

    int num_threads;
#pragma omp parallel 
    { 
    num_threads = omp_get_num_threads();
    }
    fprintf (stderr, "\n num_threads: %d\n", num_threads);

    // Data Initialization
    Key* keys;
    Key* keys_t;
    Val* vals;
    Val* vals_t;
    keys   = (Key*) _mm_malloc (size * num_threads * sizeof (Key), IDEAL_ALIGNMENT);
    keys_t = (Key*) _mm_malloc (size * num_threads * sizeof (Key), IDEAL_ALIGNMENT);
    vals   = (Val*) _mm_malloc (size * num_threads * sizeof (Val), IDEAL_ALIGNMENT);
    vals_t = (Val*) _mm_malloc (size * num_threads * sizeof (Val), IDEAL_ALIGNMENT);

    int* my_buckets;
    int* buckets;
    my_buckets = (int*) _mm_malloc ((num_buckets+3) * num_threads * sizeof (int), IDEAL_ALIGNMENT);
    buckets    = (int*) _mm_malloc ((num_buckets+3) * num_threads * sizeof (int), IDEAL_ALIGNMENT);

    for (int pass=0; pass < 3; pass++) {
        //  Initialize Uniform random data
        //  Keys and Vals
#pragma omp parallel
        {
        int tid = omp_get_thread_num();
        int num_threads = omp_get_num_threads();
        for (int i = 0; i < size; i++) {
            keys[tid*size+i] = (static_cast<long> (rand()) << 32) | rand();
//          keys[tid*size+i] = tid * size + i; 
            vals[tid*size+i] = i;
            keys_t[tid*size+i] = 0;
            vals_t[tid*size+i] = 0;
        }
    }

    // Timing variables
    double compute_time=0.;
    int timer_handle;
    bool swap = false;

    timer_handle = Timer::Start();            

    for (int shift = 0; shift < num_bits; shift += num_radix_bits) {
        long mask = static_cast<long>((1 << num_radix_bits) - 1);
#pragma omp parallel 
        {
//      double tb, tp, tw;
//      tb = 0.0;
//      tp = 0.0;
//      tw = 0.0;
//      int timer_handle;

        int tid = omp_get_thread_num();
        int num_threads = omp_get_num_threads();
        int bi = tid * (num_buckets + 3);
        int ni = tid * size;

        // Clear bucket counts
        for (int i = 0; i < num_buckets+3; i++) {
          my_buckets[bi + i] = 0;
        }

        if (swap) {
//          timer_handle = Timer::Start();            
            // Compute histogram 
            for (int i = 0; i < size; i++) {
                my_buckets[bi + ((keys_t[ni + i] >> shift) & mask)]++; 
            }
//          tb += Timer::Stop(timer_handle, "hist_time");
            #pragma omp barrier

            // Compute prefix-sums
//          timer_handle = Timer::Start();            
            int st = tid * ((num_buckets+3)/num_threads);
            int end =  (tid ==  num_threads-1) ? (num_buckets+3) : (tid+1) * ((num_buckets+3)/num_threads);
            for (int i = st; i < end; i++) {
                int last = buckets[i] = my_buckets[i];
                for (int j = 1; j < num_threads; j++) {
                    int temp = buckets[j*(num_buckets+3) + i] = last + my_buckets[j*(num_buckets+3) + i];
                    last = temp;
                }
            }
            #pragma omp barrier

            int offset = 0;  
            for (int i = 0; i < num_buckets+3; i++) {
                my_buckets[i+bi] = (buckets[i+bi] - my_buckets[i+bi]) + offset;
                offset += buckets[(num_threads-1) * (num_buckets+3) + i];
            } 
//          tp += Timer::Stop(timer_handle, "prefix_time");
            #pragma omp barrier
 
            // Write keys/vals into correct locations in temp array
//          timer_handle = Timer::Start();            
            for (int i = 0; i < size; i++) {
                int j = (keys_t[i+ni] >> shift) & mask;
                keys[my_buckets[bi+j]] = keys_t[ni+i];
                vals[my_buckets[bi+j]] = vals_t[ni+i];
                my_buckets[bi+j]++;
            }
//          tw += Timer::Stop(timer_handle, "write_time");
            swap = false;
        }
        else {
            // Compute histogram
//          timer_handle = Timer::Start();            
            for (int i = 0; i < size; i++) {
                my_buckets[bi + ((keys[ni + i] >> shift) & mask)]++; 
            }
//          tb += Timer::Stop(timer_handle, "hist_time");
            #pragma omp barrier

            // Compute prefix-sums
//          timer_handle = Timer::Start();            
            int st = tid * ((num_buckets+3)/num_threads);
            int end =  (tid ==  num_threads-1) ? (num_buckets+3) : (tid+1) * ((num_buckets+3)/num_threads);
            for (int i = st; i < end; i++) {
                int last = buckets[i] = my_buckets[i];
                for (int j = 1; j < num_threads; j++) {
                    int temp = buckets[j*(num_buckets+3) + i] = last + my_buckets[j*(num_buckets+3) + i];
                    last = temp;
                }
            }
            #pragma omp barrier

            int offset = 0;  
            for (int i = 0; i < num_buckets+3; i++) {
                my_buckets[i+bi] = (buckets[i+bi] - my_buckets[i+bi]) + offset;
                offset += buckets[(num_threads-1) * (num_buckets+3) + i];
            } 
//          tp += Timer::Stop(timer_handle, "prefix_time");
            #pragma omp barrier

            // Write keys/vals into correct locations in temp array
//          timer_handle = Timer::Start();            
            for (int i = 0; i < size; i++) {
                int j = (keys[i+ni] >> shift) & mask; 
                keys_t[my_buckets[bi+j]] = keys[ni+i];
                vals_t[my_buckets[bi+j]] = vals[ni+i];
                my_buckets[bi+j]++;
            }
//          tw += Timer::Stop(timer_handle, "write_time");
            swap = true;
        }
//        fprintf (stderr, "\ntid: %d \t Histogram Time: %lf seconds\n", tid, tb);
//        fprintf (stderr, "\ntid: %d \t Prefix Time: %lf seconds\n", tid, tp);
//        fprintf (stderr, "\ntid: %d \t Write Time: %lf seconds\n", tid, tw);
    }
    }
    compute_time += Timer::Stop(timer_handle, "compute_time");
    fprintf (stderr, "\n Compute Time: %lf seconds\n", compute_time);

    // Check if sorted
#pragma omp parallel 
    {
      int tid = omp_get_thread_num();
      int num_threads = omp_get_num_threads();
      for (int i = 0; i < size-1; i++)
        if (keys[tid*size+i] > keys[tid*size + (i+1)]) {
          fprintf (stderr, "\n Data not sorted..... failed %d...%d\n",tid, i);
          exit (1);
        }           
    }
    }
  _mm_free (keys);
  _mm_free (keys_t);
  _mm_free (vals);
  _mm_free (vals_t);
  _mm_free (my_buckets);
  _mm_free (buckets);
   
}
